from collections.abc import Iterator
from typing import Any, TypeVar, TypedDict

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, pause, resume, skip

AnySample = TypeVar('AnySample')


class StridingIteratorState(TypedDict):
    iterator_state: Any
    first_stride: bool


class StridingIterator(DatasetIterator[AnySample]):
    def __init__(self, iterator: Iterator[AnySample], stride_samples_n: int, first_stride: bool = True):
        self._iterator = iterator
        self._first_stride = first_stride
        self._stride_samples_n = stride_samples_n

    @property
    def stride_samples_n(self) -> int:
        return self._stride_samples_n

    def __next__(self) -> AnySample:
        if self._first_stride:
            self._first_stride = False
            return next(self._iterator)
        else:
            skip(self._iterator, self._stride_samples_n - 1)
            return next(self._iterator)

    def __skip__(self, samples_n: int):
        if self._first_stride:
            self._first_stride = False
            skip(self._iterator, 1 + (samples_n - 1) * self._stride_samples_n)
        else:
            skip(self._iterator, samples_n * self._stride_samples_n)

    def __pause__(self) -> StridingIteratorState:
        return StridingIteratorState(
            iterator_state=pause(self._iterator),
            first_stride=self._first_stride)


class StridingDataset(Dataset[AnySample]):
    def __init__(self, dataset: Dataset[AnySample], stride_samples_n: int):
        if not stride_samples_n >= 1:
            raise ValueError(f"Expected stride_samples_n>=1, but got stride_samples_n={stride_samples_n}.")
        self._dataset = dataset
        self._stride_samples_n = stride_samples_n

    @property
    def dataset(self) -> Dataset[AnySample]:
        return self._dataset

    @property
    def stride_samples_n(self) -> int:
        return self._stride_samples_n

    def __len__(self) -> int:
        return (len(self._dataset) + self._stride_samples_n - 1) // self._stride_samples_n

    def __getitem__(self, index: int, /) -> AnySample:
        return self._dataset[index * self._stride_samples_n]

    def __iter__(self) -> StridingIterator[AnySample]:
        return StridingIterator(
            iterator=iter(self._dataset),
            stride_samples_n=self._stride_samples_n)

    def __resume__(self, state: StridingIteratorState | None = None, /) -> StridingIterator[AnySample]:
        return StridingIterator(
            iterator=resume(self._dataset, state['iterator_state']),
            stride_samples_n=self._stride_samples_n,
            first_stride=state['first_stride'])
